import torch
import torch.nn as nn
import segmentation_models_pytorch as smp

class CustomUNet(nn.Module):
    def __init__(self, in_channels, num_classes, encoder_name='resnet18', pretrained=True, 
                decoder_attention_type='scse'):
        super(CustomUNet, self).__init__()
        self.net = smp.Unet(
            encoder_name=encoder_name,
            encoder_weights='imagenet' if pretrained else None,
            decoder_attention_type=decoder_attention_type,
            in_channels=in_channels,
            classes=num_classes
        )

    def forward(self, x, need_fp=False):
        feature = self.net.encoder(x)
        if need_fp:
            # print(f"Feature shapes: {[feat.shape for feat in feature]}")
            # fp_feature = [torch.cat((feat, nn.Dropout2d(0.5)(feat))) for feat in feature]
            fp_feature = [torch.cat((feat, nn.Dropout2d(0.5)(feat) if feat.shape[1] > 1 else feat)) for feat in feature]
            # print(f"Feature shapes after dropout: {[feat.shape for feat in fp_feature]}")
            decoder_output = self.net.decoder(*fp_feature)
            outs = self.net.segmentation_head(decoder_output)
            # print(f"Final output: {outs.shape}, {outs.chunk(2)[0].shape}, {outs.chunk(2)[1].shape}")
            return outs.chunk(2)

        decoder_output = self.net.decoder(*feature)
        mask = self.net.segmentation_head(decoder_output)
        return mask

    def get_backbone_params(self):
        return self.net.encoder.parameters()

    def get_decoder_params(self):
        return self.net.decoder.parameters()


if __name__ == '__main__':
    # 示例用法
    model = CustomUNet(in_channels=1, num_classes=1, encoder_name='resnet101', pretrained=True)
    input_tensor = torch.randn(1, 1, 128, 256)
    output = model(input_tensor)
    print(f"Final output: {output.shape}")